# coding=utf-8
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
# Model(建模)-引入决策树
from sklearn import tree
import pydotplus


def iris_train():
    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
    iris = load_iris()
    # 把数据分为测试数据和验证数据
    train_data, test_data, train_target, test_target = train_test_split(iris.data, iris.target, test_size=0.2,
                                                                        random_state=1)
    # 建立一个分类器
    clf = tree.DecisionTreeClassifier(criterion="entropy")
    # 训练集进行训练
    clf.fit(train_data, train_target)
    # 画图方法2-生成pdf文件
    dot_data = tree.export_graphviz(clf, out_file=None, feature_names=clf.feature_importances_,
                                    filled=True, rounded=True, special_characters=True)
    graph = pydotplus.graph_from_dot_data(dot_data)
    # 保存图像到pdf文件
    graph.write_pdf("treetop.pdf")


if __name__ == '__main__':

    iris_train()
    # iris = load_iris()
    # print(iris)

